from Linac import *
from mpi4py import MPI
import time as t
import storage
from globals import loadLinacInfo, meminfo

fileDir = '/home/ainswort/HOMDynamics/LatticeFiles/'
saveDir = '/home/ainswort/HOMDynamics/Output/Run1/'

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

#####Build linac from data#####



#####Create Pulse#####
numBunches = 5e4
tb = 1/352.2e6
Ib = 50e-3
Ibsig = 5e-3


#####Create HOMs#####
homfreq = 1e9
homQex = 1e7
homRoverQ = 10


process = 1
N_tot = 10
if rank==0:


    fileOut = storage.LinacDB(saveDir+'lin.h5',int(numBunches),'w')

    N = 0

    while process < size and N<N_tot:
        comm.send(N,dest=process,tag=3)
        print 'Sending Linac', N, 'to process', process
        N+=1
        process+=1
        
    recieved_processes=0
    while recieved_processes<N_tot:
        L=comm.recv(source=MPI.ANY_SOURCE,tag=1)
        fileOut.Write(L)
        process = comm.recv(source=MPI.ANY_SOURCE,tag=2) 
        recieved_processes+=1

        if N<N_tot:
            comm.send(dest=process,tag=3)
            print 'Sending Linac', N, 'to process', process
            N+=1

    for process in range(1,size):
        comm.send(-1,dest=process,tag=3)


else:
    while True:
        n = comm.recv(source=0,tag=3)
        if n== -1: break

        #####Create Linac ready to track#####

        Drifts = loadLinacInfo(fileDir+'HS_2011_11_23_Plasma-Surface_To_Target.txt','DR_124','DR_493')[0]
        Cavities = CavitySection(file=fileDir+'EllipticalSection.csv',E_final=2513e6)
        P1 = Pulse(numBunches,Cavities.cavNum, tb, Ib)
        P1.AddCurrentError(Ibsig)
        H1 = HOM(Cavities.cavNum,homfreq,homQex,homRoverQ)
        Lin = Linac(Cavities,Drifts,P1,H1)

        Lin.Track()
        comm.send(Lin,dest=0,tag=1)
        comm.send(rank,dest=0,tag=2)
        print Lin.Pulse.dE[:,-1]
        del(P1,H1,Lin)

